import Foundation

// The root is the top of the Radix Tree
public class Root {
	var children: [Edge]

	public init() {
		children = [Edge]()
	}

	// Returns the length (in number of edges) of the longest traversal down the tree.
	public func height() -> Int {
		// Base case: no children: the tree has a height of 1
		if children.count == 0 {
			return 1
		}
		// Recursion: find the max height of a root's child and return 1 + that max
		else {
			var max = 1
			for c in children {
				if c.height() > max {
					max = c.height()
				}
			}
			return 1 + max
		}
	}

	// Returns how far down in the tree a Root/Edge is.
	// A root's level is always 0.
	public func level() -> Int {
		return 0
	}

	// Prints the tree for debugging/visualization purposes.
	public func printRoot() {
		// Print the first half of the children
		if children.count > 1 {
			for c in 0...children.count/2-1 {
				children[c].printEdge()
				print("|")
			}
		} else if children.count > 0 {
			children[0].printEdge()
		}
		// Then print the root
		print("ROOT")
		// Print the second half of the children
		if children.count > 1 {
			for c in children.count/2...children.count-1 {
				children[c].printEdge()
				print("|")
			}
		}
		print()
	}
}

// Edges are what actually store the Strings in the tree
public class Edge: Root {
	var parent: Root?
	var label: String

	public init(_ label: String) {
		self.label = label
		super.init()
	}

	public override func level() -> Int {
		// Recurse up the tree incrementing level by one until the root is reached
		if parent != nil {
			return 1 + parent!.level()
		}
		// If an edge has no parent, it's level is one
		// NOTE: THIS SHOULD NEVER HAPPEN AS THE ROOT IS ALWAYS THE TOP OF THE TREE
		else {
			return 1
		}
	}

	// Erases a specific edge (and all edges below it in the tree).
	public func erase() {
		self.parent = nil
		if children.count > 0 {
			// For each child, erase it, then remove it from the children array.
			for _ in 0...children.count-1 {
				children[0].erase()
				children.remove(at: 0)
			}
		}
	}

	// Prints the tree for debugging/visualization purposes.
	public func printEdge() {
		// Print the first half of the edge's children
		if children.count > 1 {
			for c in 0...children.count/2-1 {
				children[c].printEdge()
			}
		} else if children.count > 0 {
			children[0].printEdge()
		}
		// Tab over once up to the edge's level
		for x in 1...level() {
			if x == level() {
				print("|------>", terminator: "")
			} else {
				print("|       ", terminator: "")
			}
		}
		print(label)
		// Print the second half of the edge's children
		if children.count == 0 {
			for _ in 1...level() {
				print("|       ", terminator: "")
			}
			print()
		}
		if children.count > 1 {
			for c in children.count/2...children.count-1 {
				children[c].printEdge()
			}
		}
	}
}

public class RadixTree {
	var root: Root

	public init() {
		root = Root()
	}

	// Returns the height of the tree by calling the root's height function.
	public func height() -> Int {
		return root.height() - 1
	}

	// Inserts a string into the tree.
	public func insert(_ str: String) -> Bool {
		//Account for a blank input. The empty string is already in the tree.
		if str == "" {
			return false
		}

		// searchStr is the parameter of the function
		// it will be substringed as the function traverses down the tree
		var searchStr = str

		// currEdge is the current Edge (or Root) in question
		var currEdge = root

		while true {
			var found = false

			// If the current Edge has no children then the remaining searchStr is
			// created as a child
			if currEdge.children.count == 0 {
				let newEdge = Edge(searchStr)
				currEdge.children.append(newEdge)
				newEdge.parent = currEdge
				return true
			}

			// Loop through all of the children
			for e in currEdge.children {
				// Get the shared prefix between the child in question and the
				// search string
				let shared = sharedPrefix(searchStr, e.label)
				var index  = shared.startIndex

				// If the search string is equal to the shared string,
				// the string already exists in the tree
				if searchStr == shared {
					return false
				}

				// If the child's label is equal to the shared string, you have to
				// traverse another level down the tree, so substring the search
				// string, break the loop, and run it back
				else if shared == e.label {
					currEdge = e
					var tempIndex = searchStr.startIndex
					for _ in 1...shared.count {
						tempIndex = searchStr.index(after: tempIndex)
					}
                    searchStr = String(searchStr[tempIndex...])
					found = true
					break
				}

				// If the child's label and the search string share a partial prefix,
				// then both the label and the search string need to be substringed
				// and a new branch needs to be created
				else if shared.count > 0 {
					var labelIndex = e.label.startIndex

					// Create index objects and move them to after the shared prefix
					for _ in 1...shared.count {
                        index = searchStr.index(after: index)
						labelIndex = e.label.index(after: labelIndex)
					}

					// Substring both the search string and the label from the shared prefix
                    searchStr = String(searchStr[index...])
                    e.label = String(e.label[labelIndex...])

					// Create 2 new edges and update parent/children values
					let newEdge = Edge(e.label)
					e.label = shared
					for ec in e.children {
						newEdge.children.append(ec)
					}
					newEdge.parent = e
					e.children.removeAll()
					for nec in newEdge.children {
						nec.parent = newEdge
					}
					e.children.append(newEdge)
					let newEdge2 = Edge(searchStr)
					newEdge2.parent = e
					e.children.append(newEdge2)
					return true
				}
				// If they don't share a prefix, go to next child
			}

			// If none of the children share a prefix, you have to create a new child
			if !found {
				let newEdge = Edge(searchStr)
				currEdge.children.append(newEdge)
				newEdge.parent = currEdge
				return true
			}
		}
	}

	// Tells you if a string is in the tree
	public func find(_ str: String) -> Bool {
		// A radix tree always contains the empty string
		if str == "" {
			return true
		}
		// If there are no children then the string cannot be in the tree
		else if root.children.count == 0 {
			return false
		}
		var searchStr = str
		var currEdge = root
		while true {
			var found = false
			// Loop through currEdge's children
			for c in currEdge.children {
				// First check if the search string and the child's label are equal
				// if so the string is in the tree, return true
				if searchStr == c.label {
					return true
				}

				// If that is not true, find the shared string b/t the search string
				// and the label
				let shared = sharedPrefix(searchStr, c.label)

				// If the shared string is equal to the label, update the curent node,
				// break the loop, and run it back
				if shared == c.label {
					currEdge = c
					var tempIndex = searchStr.startIndex
					for _ in 1...shared.count {
						tempIndex = searchStr.index(after: tempIndex)
					}
					searchStr = String(searchStr[tempIndex...])
					found = true
					break
				}

				// If the shared string is empty, go to the next child
				else if shared.count == 0 {
					continue
				}

				// If the shared string matches the search string, return true
				else if shared == searchStr {
					return true
				}

				// If the search string and the child's label only share some characters,
				// the string is not in the tree, return false
				else if shared[shared.startIndex] == c.label[c.label.startIndex] &&
				  shared.count < c.label.count {
				  	return false
				}
			}

			// If nothing was found, return false
			if !found {
				return false
			}
		}
	}

	// Removes a string from the tree
	public func remove(_ str: String) -> Bool {
		// Removing the empty string removes everything in the tree
		if str == "" {
			for c in root.children {
				c.erase()
				root.children.remove(at: 0)
			}
			return true
		}
		// If the tree is empty, you cannot remove anything
		else if root.children.count == 0 {
			return false
		}

		var searchStr = str
		var currEdge = root
		while true {
			var found = false
			// If currEdge has no children, then the searchStr is not in the tree
			if currEdge.children.count == 0 {
				return false
			}

			// Loop through the children
			for c in 0...currEdge.children.count-1 {
				// If the child's label matches the search string, remove that child
				// and everything below it in the tree
				if currEdge.children[c].label == searchStr {
					currEdge.children[c].erase()
					currEdge.children.remove(at: c)
					return true
				}

				// Find the shared string
				let shared = sharedPrefix(searchStr, currEdge.children[c].label)

				// If the shared string is equal to the child's string, go down a level
				if shared == currEdge.children[c].label {
					currEdge = currEdge.children[c]
					var tempIndex = searchStr.startIndex
					for _ in 1...shared.count {
						tempIndex = searchStr.index(after: tempIndex)
					}
					searchStr = String(searchStr[tempIndex...])
					found = true
					break
				}
			}

			// If there is no match, then the searchStr is not in the tree
			if !found {
				return false
			}
		}
	}

	// Prints the tree by calling the root's print function
	public func printTree() {
		root.printRoot()
	}
}

// Returns the prefix that is shared between the two input strings
// i.e. sharedPrefix("court", "coral") -> "co"
public func sharedPrefix(_ str1: String, _ str2: String) -> String {
	var temp = ""
	var c1 = str1.startIndex
	var c2 = str2.startIndex
	for _ in 0...min(str1.count-1, str2.count-1) {
		if str1[c1] == str2[c2] {
			temp.append( str1[c1] )
            c1 = str1.index(after:c1)
            c2 = str2.index(after:c2)
		} else {
			return temp
		}
	}
	return temp
}
